from typing import (
    List, 
    TypeVar, 
    Type, 
    Union, 
    Callable, 
    Tuple,
    Optional
)

import os

from langchain_core.runnables import Runnable

from pydantic import BaseModel, Field
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from langchain_deepseek import ChatDeepSeek

from langchain_core.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
    PromptTemplate,
)

from loguru import logger

# 定义BaseModel，依赖Pydantic库对于结构化数据的解析功能
_BM = TypeVar("_BM", bound=BaseModel)

# 定义LLM函数
class LLMFunction:
    def __init__(
        self, chain: Runnable, name: str, max_tries: int, variables: List[str], bm: Type[_BM]
    ):
        self.chain = chain
        self.name = name
        self.MAX_TRIES = max_tries
        self.variables = list(set(variables + ["input"]))
        self.bm = bm

    def _convert(self, text) -> _BM:
        from langchain_core.utils.json import parse_json_markdown
        from json import JSONDecodeError
        # 获得</think>标签之后的文本数据
        pos = text.find("</think>")
        if pos > 0:
            text = text[pos + 8:]

        text = text.strip()
        try:
            # Markdown 格式输出的JSON格式解析，支持部分JSON格式的解析，这个对于streaming=True时非常有用
            json = parse_json_markdown(text)
            return self.bm(**json)
        except JSONDecodeError as e:
            logger.error(f"Invalid json output: {text}")
            raise e

    def invoke(self, **kwargs):
        tries = 0

        while True:
            try:
                result = self.chain.invoke(kwargs)
                logger.debug(
                    f"Call {self.name} with kwargs{kwargs} \n result: \n {result}\n"
                )
                return result if not self.bm else self._convert(result.content)
            except Exception as e:
                tries += 1
                logger.error(
                    f"Call {self.name} {tries} {kwargs} with exception: {e} "
                )
                if tries >= self.MAX_TRIES:
                    return "failed"

class LLMGenerateJsonSchema(GenerateJsonSchema):
    def sort(
        self, value: JsonSchemaValue, parent_key: Optional[str] = None
    ) -> JsonSchemaValue:
        """No-op, we don't want to sort schema values at all."""
        return value                

class LLM:
    def __init__(self):
        self.llm = self._create_llm()

    def _create_llm(self) -> ChatDeepSeek:
        openai_api_base = os.environ["OPENAI_CHAT_ENDPOINT"]
        model = os.environ["OPENAI_CHAT_MODEL"]
        key = os.environ["OPENAI_CHAT_KEY"]

        chatopenai = ChatDeepSeek(
            temperature=0.5,
            api_base=openai_api_base,
            model=model,
            api_key=key,
            streaming=False,
            max_tokens=4000,
        )
        return chatopenai

    def _create_parse_prompt(
        self, prompts: List[str], variable_names: List[str] = []
    ) -> ChatPromptTemplate:
        prompts.append("输入如下：\n{input}")
        variable_names.append("input")

        user_prompt = "\n".join(filter(None, prompts))

        user_prompt = PromptTemplate(
            template=user_prompt,
            input_variables=variable_names,
        )

        human_message_prompt = HumanMessagePromptTemplate(prompt=user_prompt)
        system_message_prompt = SystemMessagePromptTemplate(prompt=PromptTemplate(template="你是AI小助手，擅长JSON格式处理和语义理解"))

        return ChatPromptTemplate.from_messages(
            [system_message_prompt, human_message_prompt]
        )

    def _safe_bracket(self, json: dict) -> str:
        return f"{json}".replace("{", "{{").replace("}", "}}")

    def create(
        self,
        name: str,
        func: Union[
            Callable[[], Tuple[List[str], List[str]]], Tuple[str, List[str]]
        ],
        bm: Type[_BM] = None,
    ) -> LLMFunction:
        if isinstance(func, Callable):
            lines, variables = func()
        else:
            line, variables = func
            lines = [line]

        if bm:
            desc = self._safe_bracket(bm.model_json_schema(mode='serialization', schema_generator=LLMGenerateJsonSchema))
            lines.append(f"""
请结合前面的要求，严格输出JSON格式内容，JSON输出的Schema定义如下：
{desc}

请理解上诉JSON格式定义，仅输出最终的JSON格式。
其它要求：
1. 不要输出推理过程
2. 不要自行补充或者臆造内容
3. 输出的JSON的内容用双引号("")，不要用单引号('')，并注意转移字符的使用
""")

        template = self._create_parse_prompt(
            prompts=lines, variable_names=variables
        )

        return LLMFunction(
            template | self.llm,
            name,
            max_tries=int(os.environ.get("MAX_TRIES", 3)),
            variables=variables,
            bm=bm
        )

class NER(BaseModel):
    """
    Example:
    {
        "entity": "王小帅", 
        "description":"王小帅独自驾驶着出租车游荡在城市空旷的街道上",
        "occurred_on":"2023/01/09"
    }
    """

    entity: str = Field(
        description="实体名称，必须是具体的人名",
        nullable = False
    )
    description: str = Field(
        description="仿照新闻报道口吻，用一句话描述日期，地点，人物，事件等信息",
        nullable = False
    )
    occurred_on: str = Field(
        description="新闻事件发生的日期，格式为: yyyy/mm/dd， 比如:2020/02/09",
        nullable = False
    )

def prompt_1():
    return [
        """给你一段新闻片段，请提取出*一条*与{entity}相关的信息并输出，不要做任何内容的延展、补充或者编撰。""",
        """输出要求：用新闻媒体表述方式，一句话描述事件，包含具体人物名称，地点，事件描述等""",
        """其它要求：""",
        """1.entity字段：必须是具体的人名，不可以是代名词比如：他/她，不可以是职位如：市委书记"""
        """2.description字段：不需要加入时间日期，不要自行扩充新闻信息，严格尊重原文内容的意思""",
        """3.occurred_on字段：假设今天的日期是({date})，结合新闻片段中时间的描述，*仔细思考*事件发生的日期，不得使用昨天/上周日等相对日期，""",
    ],["date", "entity"]

def test_json_output():
    llm = LLM().create("extract", prompt_1, NER)
    a = llm.invoke(
        entity="领导", 
        date="2025/02/02", 
        input="""昨日，市委书记陈小明带队赴北京拜访对接央企国企，深入学习贯彻党的二十届三中全会精神和习近平总书记对江苏工作重要讲话精神"""
    )
    assert isinstance(a, NER)
    assert "陈小明" in a.entity
    assert "2025/02/01" == a.occurred_on
    b = llm.invoke(
        entity="人物", 
        date="2025/02/05", 
        input="""今天早上，市民陈女士因为最近家中要办酒席，一次性购买了三筐砂糖橘，每筐售价95元。作为冬季“明星”水果，砂糖橘备受市民青睐。"""
    )
    
    assert isinstance(b, NER)
    assert "陈女士" in b.entity
    assert "2025/02/05" == b.occurred_on

def test_benchmark():
    test_prompts = [
        "请为我总结一下关于全球气候变化的最新新闻。",
        "太阳系中最远的行星是什么？",
        "写一首关于冬季的诗歌，描述雪花飘落的美丽景象。",
        "如何使用Python连接到MySQL数据库并执行查询？",
        "我最近感到非常失落，工作压力很大，不知道该如何排解情绪。",
        "推荐我几本适合初学者阅读的机器学习书籍。",
        "计算积分 ∫(x^2 + 3x + 2) dx 从 0 到 1 的值。",
        "将这段话从中文翻译成英文：‘我喜欢编程，它能让我解决复杂的问题。’",
        "描述一下第二次世界大战的主要原因和影响。",
        "为我们的最新智能手机写一个简短的市场宣传文案，突出其创新功能。"
    ]
    import time
    llm = LLM().create("extract", ("请仔细思考并回答",[]))

    for idx, prompt in enumerate(test_prompts):
        # 计时器开始
        start_time = time.time()

        # 请求模型生成文本
        response = llm.invoke(input=prompt)

        # 计算延迟 (Latency)
        latency = time.time() - start_time

        # 计算 token 速率
        tokens_used = response.response_metadata["token_usage"]["completion_tokens"]  # 根据生成的文本计算 token 数量
        token_rate = tokens_used / latency  # 每秒处理的 token 数量

        # 打印每个测试文本的结果
        print(f"\nTest {idx + 1}: {prompt}")
        # print(f"Response: {response[:100]}...")  # 只打印前100个字符以避免过长输出
        print(f"Latency: {latency:.4f} seconds \t")
        print(f"Token Rate: {token_rate:.2f} tokens per second")
    